import torch as t
import torchvision as tv
from PIL import Image

# 分类数组
classes = {
    'tomatomodel': ['靶点病', '黄曲叶病毒', '菌斑症', '早疫病',
            '健康', '晚疫病', '叶霉症', '七星叶斑病', '蜘蛛螨及二斑蜘蛛螨'],
    'applemodel': ['赤霉病', '黑腐病', '苹果锈', '健康'],
    'cherrymodel': ['白粉病', '健康'],
    'cornmodel': ['灰斑病', '锈病', '北方叶枯病', '健康'],
    'grapemodel': ['黑腐病', '葡萄黑死病(ESCA)', '叶枯病', '健康'],
    'peppermodel': ['菌斑症', '健康'],
    'potatomodel': ['早疫病', '晚疫病', '健康'],
    'peachmodel': ['菌斑症', '健康']
}

class ResnetClassifier:
    def __init__(self):
        bCuda = t.cuda.is_available()  # 是否开启 GPU
        bCuda = False  # 不启用GPU
        self.device = t.device("cuda:0" if bCuda else "cpu")

        img_size = 256  # 图片大小，可以改

        # 对Tensor进行变换 颜色转换   mean=给定均值：(R,G,B) std=方差：（R，G，B）
        normalize = tv.transforms.Normalize(
            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        self.transform = tv.transforms.Compose(
            [tv.transforms.Resize([img_size, img_size]), tv.transforms.CenterCrop([img_size, img_size]),
            tv.transforms.ToTensor(), normalize])

    def loadModel(self, name):
        # 直接加载
        model = t.load(f'resource/{name}.pkl', map_location=t.device('cpu'))
        self.model = model.to(self.device)  # GPU
        self.model.eval()  # 运行模式
        
    def classify(self, modelName, image_PIL):
        self.loadModel(modelName)
        

        t.no_grad()
        
        # imshow(image_PIL)

        image_tensor = self.transform(image_PIL)
        # 以下语句等效于 img = torch.unsqueeze(image_tensor, 0)
        image_tensor.unsqueeze_(0)
        # 没有这句话会报错
        image_tensor = image_tensor.to(self.device)
        out = self.model(image_tensor)
        # 得到预测结果，并且从大到小排序
        _, indices = t.sort(out, descending=True)

        # 返回每个预测值的百分数
        percentage = t.nn.functional.softmax(out, dim=1)[0] * 100

        # 是否显示每个分类的预测值
        item = indices[0]
        # if isShowSoftmax:
        #     for idx in item:
        #         ss = percentage[idx]
        #         value = ss.item()
        #         name = classes[idx]
        #         print('名称：', name, '预测值：', value)

        # 预测最大值
        _, predicted = t.max(out.data, 1)
        maxPredicted = classes[modelName][predicted.item()]
        maxAccuracy = percentage[item[0]].item()
        return maxPredicted, maxAccuracy

if __name__ == '__main__':
    classifier = ResnetClassifier()
    print(classifier.classify('example.jpg'))
